from __future__ import annotations

import base64
import os
from typing import Optional

from flask import Flask, jsonify, request, send_from_directory
from langchain_openai import ChatOpenAI


def file_to_data_url(filename: str, data: bytes) -> str:
  ext = os.path.splitext(filename)[1].lower()
  mime = {
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png": "image/png",
    ".gif": "image/gif",
    ".webp": "image/webp",
    ".bmp": "image/bmp",
    ".tiff": "image/tiff",
    ".svg": "image/svg+xml",
  }.get(ext, "application/octet-stream")
  b64 = base64.b64encode(data).decode("utf-8")
  return f"data:{mime};base64,{b64}"


def build_content(prompt: str, image_url: Optional[str], data_url: Optional[str]):
  parts = []
  if prompt:
    parts.append({"type": "text", "text": prompt})
  if image_url:
    parts.append({"type": "input_image", "image_url": image_url, "detail": "high"})
  elif data_url:
    parts.append({"type": "input_image", "image_url": data_url, "detail": "high"})
  return parts


app = Flask(__name__, static_folder="static", static_url_path="/static")


@app.route("/")
def index():
  return send_from_directory("static", "index.html")


@app.post("/recognize")
def recognize():
  try:
    prompt = request.form.get("prompt", "请识别图片中的文字并总结要点。")
    image_url = request.form.get("image_url")
    base64_data = request.form.get("data_url")

    if not image_url and not base64_data and "file" in request.files:
      f = request.files["file"]
      base64_data = file_to_data_url(f.filename, f.read())

    api_key = os.environ.get("DEEPSEEK_API_KEY") or request.form.get("api_key")
    base = request.form.get("base") or os.environ.get("DEEPSEEK_API_BASE", "https://api.deepseek.com/v1")
    model = request.form.get("model") or os.environ.get("DEEPSEEK_VISION_MODEL", "deepseek-vl")
    if not api_key:
      return jsonify({"error": "Missing DEEPSEEK_API_KEY"}), 400

    llm = ChatOpenAI(api_key=api_key, base_url=base, model=model, temperature=0.2, max_tokens=1024)
    content = build_content(prompt, image_url, base64_data)
    if not content:
      return jsonify({"error": "Provide image_url, data_url or file"}), 400

    resp = llm.invoke([{ "role": "user", "content": content }])
    text = resp.content if hasattr(resp, "content") else str(resp)
    return jsonify({"text": text})
  except Exception as e:
    return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
  app.run(host="0.0.0.0", port=8000, debug=True)


